#!/usr/bin/python
# coding:utf-8

import time
from typing import Callable

import orjson
from fastapi import Request, Response
from fastapi.routing import APIRoute

from ..file_logger import logger

"""
自定义APIRoute, 用于记录request和response
"""


class TimedRoute(APIRoute):
    def get_route_handler(self) -> Callable:
        original_route_handler = super().get_route_handler()

        # https://fastapi.tiangolo.com/advanced/custom-request-and-route/
        # https://juejin.cn/post/6972031031155097631
        # https://stackoverflow.com/questions/60098005/fastapi-starlette-get-client-real-ip
        # https://gist.github.com/philippegirard/7cdbec8036561285b5579d8d334b20ba
        # https://stackoverflow.com/questions/63400683/python-logging-with-loguru-log-request-params-on-fastapi-app

        async def custom_route_handler(request: Request) -> Response:
            before = time.time()
            response: Response = await original_route_handler(request)
            duration = (time.time() - before) * 1000

            # https://stackoverflow.com/questions/64115628/get-starlette-request-body-in-the-middleware-context
            request_body = (await request.body()).decode()
            if request_body:
                request_data = orjson.loads(request_body)
            else:
                request_data = None
            response_body = response.body.decode()
            if response_body:
                response_data = orjson.loads(response_body)
            else:
                response_data = None
            logger.info(
                {
                    "op": "RecordRouteRequestAndResponse",
                    "requestMethod": request.method,
                    "url": str(request.url),
                    "headers": request.headers.items(),
                    "ip": request.client.host,
                    "request": request_data,
                    "response": response_data,
                    "routeDuration(ms)": f"{duration:.3f}",
                }
            )
            return response

        return custom_route_handler
